#if _MSC_VER
 #pragma once
#endif

#ifndef __HERZELEID_STORAGE_INCLUDED__
 #define __HERZELEID_STORAGE_INCLUDED__
 #include "HerzeleidStorageReservedKeys.hxx"

 #define HerzeleidUtilityNamespaceBegin namespace Herzeleid { namespace Utility { 
 #define HerzeleidUtilityNamespaceImpl Herzeleid::Utility 
 #define HerzeleidUtilityNamespaceEnd ; }; } 

 #ifdef HERZELEID_DESKTOP
  #define ref 
  #define internal public
 #endif // HERZELEID_DESKTOP

HerzeleidUtilityNamespaceBegin template <class TPriorityKey = UInt32, class TKeyComparator = std::less<UInt32>> 
// Storage manager ensures the attached resources are removed in a correct order (for more info StorageReservedKeys) when destroyed, 
// relying on the provided to the Storage<...>::DemandAttach(...) function keys and their comparator (class 2nd type param).
ref class Storage : public Standard::Named
{
protected:
	typedef std::mutex Mutex;
	// Holds the resource pointer, host`s collection and a bound key. 
	struct WrappedData;
	// Holds the ComPtr pointer
	struct ComWrappedData;
	// Holds the shared pointer
	template <class TObject>
	struct AutoWrappedData;
	// Manages the keyed resource collection. 
	struct ResourceMap;

internal:
	// Internally used: manages the keyed resource collection.
	typedef std::multimap<TPriorityKey, std::shared_ptr<WrappedData>, TKeyComparator> KeyedMultimap;

	Storage() : m_Multimap(std::make_shared<ResourceMap>()), Named()
	{
		std::lock_guard<std::mutex> lGuard(m_Multimap->m_Mtx);

#ifdef HERZELEID_DEBUG // Avoid noops
		{
			Char Buffer[16];
			swprintf_s( 
				Buffer, L"%s%d", L"Pool", 
				sm_uNumInstancesAllocated++ 
				);

			DebugName = Buffer;
		}
#endif

		m_Multimap.get()->m_HostPool = this;
	}

	// Set debug name and allocate resource multimap 
	Storage(const String16 &szName) : 
		m_Multimap(std::make_shared<ResourceMap>())
#ifdef HERZELEID_DEBUG
		, Named(szName) // if debugging, set the name 
#else
		, Named() // otherwise, noop
#endif
	{
		std::lock_guard<std::mutex> lGuard(m_Multimap->m_Mtx);
		m_Multimap.get()->m_HostPool = this;
	}

	virtual ~Storage()
	{
		std::lock_guard<std::mutex> lGuard(m_Multimap->m_Mtx);
		m_Multimap.get()->clear();
	}

	KeyedMultimap const *const GetMultimapPtr() const 
	{
		return(m_Multimap.get());
	}

	template <class TObject>
	// Scenario: 
	// 1> Checks whether the allocated pointer is provided (if it is not, generates E_POINTER).
	// 2> Checks if there are any expired resources with the same key (erases if any).
	// 3> Adds the provided data to the collection. 
	// \note Does not make any checks whether the key is correct. Consider this when using the default library resource pool (not to replace the allocated COMs). 
	HResultConst DemandAttach( _In_ TPriorityKey Key, _In_ TObject *DataPtr )
	{
		// The resource should be allocated (do not store empty resources). 
		if ( !DataPtr ) return(E_POINTER);

		std::lock_guard<std::mutex> lGuard(m_Multimap->m_Mtx);

		// Get the range of the elements with the same key 
		auto Range = m_Multimap->equal_range(Key);
		// Iterate through the queued elements and make a cleanup if they had expired, otherwise noop. 
		for ( auto Iterator = Range.first; Iterator != Range.second; Iterator++ )
		{
			// Remove all the released resources 
			if (!Iterator->second.get())
			{
				m_Multimap->erase(Iterator);
			}
		}

		WrappedData *lpData;
		{
			// The resource should be a COM resource (or implement ref counting at least)
			auto lComPtr = dynamic_cast<IUnknown*>(DataPtr);
			if ( !lComPtr ) lpData = new AutoWrappedData<TObject>(Key, m_Multimap, DataPtr);
			else lpData = new ComWrappedData(Key, m_Multimap, lComPtr);
		}

		std::shared_ptr<WrappedData> Data(lpData);
		m_Multimap->insert(std::make_pair(Key, Data));

		return(S_OK);
	}

protected:

	// Internally used: manages the keyed resource collection simultaneously
	typedef struct ResourceMap : public KeyedMultimap 
	{
		using KeyedMultimap::equal_range;
		using KeyedMultimap::insert;
		using KeyedMultimap::clear;
		using KeyedMultimap::find;
		// Pool that maintains the data 
		Standard::Named *m_HostPool;
		Mutex m_Mtx;
	}; // ResourceMap

	typedef std::shared_ptr<ResourceMap> SharedResourceMapPtr;

	// Internally used: holds the resource pointer, host collection and a bound key. 
	struct WrappedData // : protected Logging::Broadcaster
	{
		SharedResourceMapPtr m_BoundMap;
		TPriorityKey m_BoundKey;

	public:
		WrappedData(TPriorityKey Key, const SharedResourceMapPtr &Map ) : 
			m_BoundKey(Key), m_BoundMap(Map)
		{
			/*GetAir()->QueueMessageWithArgs(
				"WrappedData::WrappedData: Pool:%s Resource created\n", 
				m_BoundMap.get()->m_HostPool->GetDebugName().c_str()
				);*/
		}

		virtual ~WrappedData()
		{
			std::lock_guard<std::mutex> lock(m_BoundMap->m_Mtx);

			// Here to check for the data expiry. The DemandCreate can run on a different thread 
			// at the same time as a previous instance is being destroyed. 
			auto Range = m_BoundMap->equal_range(m_BoundKey);
			for ( auto Iterator = Range.first; Iterator != Range.second; Iterator++ )
			{
				// The replacement objects just been added must not be erased. 
				if (!Iterator->second.get())
				{
					// But all the expired resources should be cleaned up. 
					m_BoundMap->erase(Iterator);
				}
			}

			/*GetAir()->QueueMessageWithArgs(
				"WrappedData::~WrappedData: Pool:%s Resource removed\n", 
				m_BoundMap.get()->m_HostPool->GetDebugName().c_str()
				);*/
		}

		virtual bool Expired(
			) pure;
	}; // WrappedData

	class ComWrappedData : public WrappedData
	{
		Microsoft::WRL::ComPtr<IUnknown> m_BoundData;

	public:
		ComWrappedData( _In_ TPriorityKey Key, _In_ const SharedResourceMapPtr &Map, _In_ IUnknown *DataPtr ) : 
			WrappedData(Key, Map)
		{
			m_BoundData.Attach(DataPtr);
		}

		virtual ~ComWrappedData()
		{
			/*GetAir()->QueueMessageWithArgs(
				"ComWrappedData::~ComWrappedData: Com Resource removed\n"
				);*/
		}

		virtual bool Expired() override
		{
			return (m_BoundData.Get() ? true : false);
		}
	};

	template<class TObject>
	class AutoWrappedData : public WrappedData
	{
		std::shared_ptr<TObject> m_BoundData;

	public:
		AutoWrappedData( _In_ TPriorityKey Key,  _In_ const SharedResourceMapPtr &Map, _In_ TObject *DataPtr ) : 
			WrappedData(Key, Map), m_BoundData(DataPtr)
		{
		}

		virtual ~AutoWrappedData()
		{
			/*GetAir()->QueueMessageWithArgs(
				"AutoWrappedData::~AutoWrappedData: Generic Resource removed\n"
				);*/
		}

		virtual bool Expired() override
		{
			std::lock_guard<std::mutex> lock(m_BoundMap->m_Mtx);
			return (m_BoundData.get() ? true : false);
		}
	};

private:
	// Resource storage (managed internally). 
	std::shared_ptr<ResourceMap> m_Multimap;

#ifdef HERZELEID_DEBUG // Its not needed if NDEBUG
	static UInt32 sm_uNumInstancesAllocated;
#endif

	// Prevent copying.
	Storage( Storage const& );
	Storage &operator=( Storage const& );

}; 

template<class TPriorityKey, class TKeyComparator>
UInt32 Storage<TPriorityKey, TKeyComparator>::sm_uNumInstancesAllocated = 0u;

HerzeleidUtilityNamespaceEnd; 

 #ifdef HERZELEID_DESKTOP
  #undef ref 
  #undef internal
 #endif // HERZELEID_DESKTOP

#endif // !HERZELEID_STORAGE_INCLUDED__
